import pandas as pd

df_NH = pd.read_csv('static/data/pretreatment_NH_data.csv')
df_SHH = pd.read_csv('static/data/pretreatment_SHH_data.csv')
df_regression_line = pd.read_csv('static/data/liner_Regression.csv')


def data_heatmap():
    range_list = []
    for num in range(1, 21):
        range_list.append(df_NH[['region', 'unit_price']].groupby(['region']).apply(
            lambda x: x[(x["unit_price"] <= 1000 * num) & (1000 * (num - 1) < x["unit_price"])].count())[
                              'unit_price'].rename('{}k-{}k'.format(num, num - 1)))
    df = pd.DataFrame(range_list)
    data = [[i,j,int(df.iloc[j,i])] for j in range(len(df.index)) for i in range(len(df.columns))]
    region = df.columns.tolist()
    price_range = df.index.tolist()
    return {
        'data': data,
        'region': region,
        'price_range': price_range
    }


def data_graph():
    df_graph = df_NH.drop('house_type', axis=1).join(
        df_NH['house_type'].str.split(',', expand=True).stack().reset_index(level=1, drop=True).rename('house_type'))
    df_edges = df_graph[['house_type', 'name']].groupby(['house_type', 'name']).count().index.tolist()
    edges = [{'source': i[1], 'target': i[0]} for i in df_edges]
    df_data_type = df_graph[['house_type']].groupby(['house_type']).count().index.tolist()
    data_type = [{'name': i, 'category': 1} for i in df_data_type]
    df_data_house = df_graph[['name']].groupby(['name']).count().index.tolist()
    data_house = [{'name': i, 'category': 0} for i in df_data_house]
    data_house.extend(data_type)
    return {
        'data': data_house,
        'edges': edges,
        'category': [
            {'name': '房产商'},
            {'name': '户型'}
        ]
    }


def data_scatter():
    line_data = df_regression_line.values.tolist()
    scatter_data = df_SHH[['unit_price', 'area(㎡)']].values.tolist()
    return {
        'line_data': line_data,
        'scatter_data': scatter_data
    }


def data_pie_and_annular():
    pie_data = [{'value': df_NH.count().values.tolist()[0], 'name': '新房'},
                {'value': df_SHH.count().values.tolist()[0], 'name': '二手房'}]
    df_NH_annular = df_NH[['region', 'name']].groupby(['region']).count()
    df_SHH_annular = df_SHH[['region', 'name']].groupby(['region']).count()
    annular_data = df_SHH_annular.add(df_NH_annular, fill_value=0).reset_index().values.tolist()
    annular_data = [{'value': i[1], 'name': i[0]} for i in annular_data]
    return {
        'pie_data': pie_data,
        'annular_data': annular_data
    }


def data_line_and_bar():
    line_data = df_SHH.groupby(['orientation']).agg({'unit_price': ['mean']}).reset_index().values[:, 1].tolist()
    name = df_SHH.groupby(['orientation']).agg({'unit_price': ['mean']}).reset_index().values[:, 0].tolist()
    bar_data = df_SHH.groupby(['orientation']).count().reset_index().values[:, 1].tolist()
    return {
        'line_data': line_data,
        'bar_data': bar_data,
        'name': name
    }


def data_wordcloud():
    df_NH_wordcloud = df_NH.drop('tag', axis=1).join(
        df_NH['tag'].str.split(',', expand=True).stack().reset_index(level=1, drop=True).rename('tag'))
    df_SHH_wordcloud = df_SHH.drop('tag', axis=1).join(
        df_SHH['tag'].str.split(',', expand=True).stack().reset_index(level=1, drop=True).rename('tag'))
    df_NH_wordcloud = df_NH_wordcloud[df_NH_wordcloud['tag'] != '无']
    df_SHH_wordcloud = df_SHH_wordcloud[df_SHH_wordcloud['tag'] != '无']
    wordcloud_NH_data = df_NH_wordcloud.groupby(['tag']).count().reset_index().values[:, :2].tolist()
    wordcloud_SHH_data = df_SHH_wordcloud.groupby(['tag']).count().reset_index().values[:, :2].tolist()
    wordcloud_data = wordcloud_NH_data + wordcloud_SHH_data
    wordcloud_data = [{'name': i[0], 'value': i[1]} for i in wordcloud_data]
    return {
        'data':wordcloud_data
    }